import numpy as np
from enum import Enum
import unittest

from sqlalchemy import true

EMPTY_PIECES = 0
WHITE_PIECES = 1
BLACK_PIECES = 2

GOBANG_ROW = 15
GOBANG_COL = 15
GOBANG_GROUP = 5


class GobangEnv(object):
    chessboard = []
    status = 0

    # row, col： 棋盘规格；
    # group：group连成一串则赢
    def __init__(self, row, col, group):
        self.row = row
        self.col = col
        self.group = group
        self.chessboard = np.zeros((self.row, self.col), dtype=int)
        self.status = self.statusFromBoard(self.chessboard)

    def reset(self):
        self.chessboard = np.zeros((self.row, self.col), dtype=int)
        self.status = self.statusFromBoard(self.chessboard)
        return self.status

    def reset_with_status(self, status):
        self.chessboard = self.boardFromStatus(status)
        self.status = status
        return self.status

    def step(self, action):
        # 首先判断动作的合法性
        valid = self.isValidAction(action)
        if (not valid):
            print("x:%d,y:%d,color:%d" % (action[0], action[1], action[0]))
            raise Exception("unexpected action")
        x = action[0]
        y = action[1]
        self.chessboard[x, y] = action[2]
        self.status = self.statusFromBoard(self.chessboard)
        reward = self.calculateReward(self.chessboard)
        done = True if (reward == 100 or reward == -100) else False
        if (not done):
            hasEmpty = False
            for i in range(self.row):
                for j in range(self.col):
                    if (self.chessboard[i, j] == 0):
                        hasEmpty = True
                        break
                if (not hasEmpty):
                    break

            done = True if (not hasEmpty) else False

        return self.status, reward, done, 0

    def calculateReward(self, chessboard):
        reward = 0
        visited = np.zeros((self.row, self.col), dtype=int)
        result = False
        for i in range(self.row):
            for j in range(self.col):
                if ((visited[i][j] == 1) or (chessboard[i][j] == 0)):
                    continue

                color = chessboard[i, j]
                result = self.dfs(chessboard, visited, (i, j), -1, 1)
                if (color == WHITE_PIECES and result):
                    reward = 100
                    return reward

                if (color == BLACK_PIECES and result):
                    reward = -100
                    return reward

        return reward

    def dfs(self, chessboard, visited, position, direct, len):
        x = position[0]
        y = position[1]
        visited[x, y] = 1
        color = chessboard[x, y]
        nextPosVec = [(x-1, y+1), (x, y+1), ((x+1), y+1), (x+1, y)]

        if (direct == -1):
            direct = 0
        orderList = [direct]
        for order in [0, 1, 2, 3]:
            if (order != direct):
                orderList.append(order)

        nextPos = nextPosVec[direct]
        result = False

        for order in orderList:
            nextPos = nextPosVec[order]
            # 检查位置是否合法
            if (not (0 <= nextPos[0] and nextPos[0] < self.col
                     and 0 <= nextPos[1] and nextPos[1] < self.row)):
                continue

            # 检查棋子颜色是否与前一个棋子相同，如果不同，则停止搜索本节点
            nextColor = chessboard[nextPos[0], nextPos[1]]
            if (nextColor != color):
                continue

            # 检查是否一串棋子，如果是则len+1，然后判断串起来是否达到5个，达到5个则dfs退出，reward计算过程结束。
            # 如果不是一串棋子，则最近的两个节点组成一串重新搜索
            if (order == direct):
                if (len+1 == self.group):
                    return True

                result = self.dfs(chessboard, visited,
                                  nextPos, order, len+1)
            else:
                result = self.dfs(chessboard, visited, nextPos, order, 2)

            if(result):
                return True
        return False

    def boardFromStatus(self, status):
        board = np.zeros((self.row, self.col), dtype=int)
        print(status)
        for i in range(self.row):
            for j in range(self.col):
                temp = int(pow(3, i*(self.col)+j+1))
                remainder = status % temp
                if (not (i == 0 and j == 0)):
                    board[i, j] = int(
                        remainder / int(pow(3, i*(self.col)+j)))
                # print("i:%d,j:%d,board[i,j]:%d,temp:%d" %
                #       (i, j, board[i, j], temp))

        return board

    def statusFromBoard(self, board):
        status = int(0)
        for i in range(self.row):
            for j in range(self.col):
                temp = board[i, j] * pow(3, i*(self.col)+j)
                status = status + int(temp)

        return int(status)

    def isValidAction(self, action):
        x = action[0]
        y = action[1]
        if (not (0 <= x and x < self.row and 0 <= y and y < self.col)):
            action.render()
            return False

        if (self.chessboard[x, y] != 0):
            return False

        return True

    def act(self, action):
        valid = self.isValidAction(action)
        if (not valid):
            action.render()
            raise Exception("unexpected action")
        x = action.x
        y = action.y
        self.chessboard[x, y] = action[2]
        self.status = self.statusFromBoard(self.chessboard)

    def render(self):
        print(self.chessboard)


class TestGoBang(unittest.TestCase):

    def test_upper(self):
        self.assertEqual('foo'.upper(), 'FOO')

    def test_isupper(self):
        self.assertTrue('FOO'.isupper())
        self.assertFalse('Foo'.isupper())

    def test_split(self):
        s = 'hello world'
        self.assertEqual(s.split(), ['hello', 'world'])
        # check that s.split fails when the separator is not a string
        with self.assertRaises(TypeError):
            s.split(2)

    def test_gobang(self):
        row = 10
        col = 10
        gobang = GobangEnv(row, col, 3)

        next_obs, reward, done, test = gobang.step([1, 0, WHITE_PIECES])
        # gobang.render()
        next_obs, reward, done, test = gobang.step([1, 1, WHITE_PIECES])
        # gobang.render()
        next_obs, reward, done, test = gobang.step([1, 2, WHITE_PIECES])
        next_obs, reward, done, test = gobang.step([1, 3, WHITE_PIECES])
        next_obs, reward, done, test = gobang.step([1, 4, WHITE_PIECES])
        # gobang.render()
        self.assertEqual(reward, 100)

        gobang.reset()
        next_obs, reward, done, test = gobang.step([1, 0, BLACK_PIECES])
        next_obs, reward, done, test = gobang.step([2, 1, BLACK_PIECES])
        next_obs, reward, done, test = gobang.step([3, 2, BLACK_PIECES])
        next_obs, reward, done, test = gobang.step([4, 3, BLACK_PIECES])
        next_obs, reward, done, test = gobang.step([5, 4, BLACK_PIECES])
        # gobang.render()
        self.assertEqual(reward, -100)

    def test_board_status(self):
        row = 10
        col = 10
        env = GobangEnv(row, col, 3)

        status = 1
        board = env.boardFromStatus(status)
        # print(board)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)

        status = 2
        board = env.boardFromStatus(status)
        # print(board)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)

        status = 500
        board = env.boardFromStatus(status)
        # print(board)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)

        status = 1899
        board = env.boardFromStatus(status)
        # print(board)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)

        status = 3600
        board = env.boardFromStatus(status)
        # print(board)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)
        return

    def test_board_status2(self):
        env = GobangEnv(15, 15, 5)
        board_list = [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
        board = np.array(board_list)

        status = env.statusFromBoard(board)
        board_new = env.boardFromStatus(status)
        status_new = env.statusFromBoard(board)
        self.assertEqual(status, status_new)

        return


if __name__ == '__main__':
    unittest.main()
